import idc
from example.ida.utils import guess_machine,expr2colorstr
from future.utils import viewitems, viewvalues

from miasm.arch.x86.regs import *

from miasm.ir.ir import IRBlock, AssignBlock
from miasm.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAssign,ExprCond, ExprId, ExprMem,ExprSlice

class BaseSolver:
    def __init__(self,rule_name,tag='') -> None:
        self.rule_name=rule_name
        self.tag=tag
        self.action=None
        pass
    def lift(self):
        machine=self.machine
        loc_db=self.loc_db
        asmcfg=self.asmcfg
        # lifter
        lifter_model_call=machine.lifter_model_call
        use_ida_stack=True
        class IRADelModCallStack(lifter_model_call):
            def call_effects(self, addr, instr):
                assignblks, extra = super(IRADelModCallStack, self).call_effects(addr, instr)
                if use_ida_stack:
                    stk_before = idc.get_spd(instr.offset)
                    stk_after = idc.get_spd(instr.offset + instr.l)
                    stk_diff = stk_after - stk_before
                    # print(hex(stk_diff))
                    call_assignblk = AssignBlock(
                        [
                            ExprAssign(self.ret_reg, ExprOp('call_func_ret', addr)),
                            ExprAssign(self.sp, self.sp + ExprInt(stk_diff, self.sp.size))
                        ],
                        instr
                    )
                    return [call_assignblk], []
                else:
                    if not dontmodstack:
                        return assignblks, extra
                    out = []
                    for assignblk in assignblks:
                        dct = dict(assignblk)
                        dct = {
                            dst:src for (dst, src) in viewitems(dct) if dst != self.sp
                        }
                        out.append(AssignBlock(dct, assignblk.instr))
                return out, extra
        lifter=IRADelModCallStack(loc_db)
        # main_ircfg
        main_ircfg=lifter.new_ircfg_from_asmcfg(asmcfg)
        self.lifter=lifter
        self.ircfg=main_ircfg
        pass
    def Solve(self,obj):
        self.asmcfg=obj['asmcfg']
        self.head=obj['head']
        self.mdis=obj['dis']
        self.loc_db=obj['loc_db']
        self.func_now=obj['func_now']
        self.machine=obj['machine']
        self.addr=obj['addr']
        print('[DEBUG]Analygis start with <Rule>%s,%s' % (self.rule_name,self.tag))
    
    def format_action(self):
        if self.action is None:
            return
        str_add_blk=self.action.str_add_blk()
        str_del_blk=self.action.str_del_blk()
        str_add_edge=self.action.str_add_edge()
        str_del_edge=self.action.str_del_edge()
        return str_add_blk,str_del_blk,str_add_edge,str_del_edge
        # print(str_add_blk,str_del_blk,str_add_edge,str_del_edge)

if __name__=='__main__':
    base=BaseSolver('test')
    base.Solve({})